from flask import jsonify, request, g
from app.errors import Success
import datetime
from uuid import uuid4
from app.libs.redprint import RedPrint
from app.log.logs import (
    CProfileTemplateLog,
    UProfileTemplateLog,
    DProfileTemplateLog,
)
from app.db.models.vulndb import (
    SystemKeywords,
    CustomKeywords,
    WhiteListKeywords,
    SystemProfileTemplate,
    CustomProfileTemplate,
    KbVuln,
    NameType
)
from app.validators import kb as kb_form

api = RedPrint('kb')


@api.route('/keywords', methods=['GET'])
def keywords():
    form = kb_form.QueryKeywordForm().validate_for_api()
    query = {'is_deleted': False}
    if word := form.word.data:
        query['word__icontains'] = word
    system_keywords = SystemKeywords.objects.filter(**query)
    custom_keywords = CustomKeywords.objects.filter(custom_id=str(form.uid.data), **query)

    system_keywords_pipeline = [
        # {'$group': {'_id': '$type', 'type': {'$first': '$type'}, 'system_keywords_count': {'$sum': 1},'system_keywords': {'$push': '$word'}}},
        {'$group': {'_id': '$type', 'type': {'$first': '$type'}, 'system_keywords_count': {'$sum': 1}}},
        {'$project': {'_id': 0}}]
    custom_keywords_pipeline = [
        {'$group': {'_id': '$type', 'type': {'$first': '$type'}, 'custom_keywords_count': {'$sum': 1}}},
        {'$project': {'_id': 0}}]

    resp = {}
    for item in system_keywords.aggregate(system_keywords_pipeline):
        resp.setdefault(item['type'], {}).update(item)

    for item in custom_keywords.aggregate(custom_keywords_pipeline):
        resp.setdefault(item['type'], {}).update(item)

    return jsonify(sorted(resp.values(), key=lambda x: x['type'], reverse=False))


@api.route('/keywords/<int:_type>', methods=['GET'])
def keywords_type(_type):
    form = kb_form.QueryKeywordTypeForm().validate_for_api()
    return jsonify({
        'system_keywords': list(SystemKeywords.objects.filter(type=_type, is_deleted=False).scalar('word')),
        'custom_keywords': list(CustomKeywords.objects.filter(custom_id=str(form.uid.data), type=_type).scalar('word'))
    })


@api.route('/custom_keywords', methods=['POST'])
def save_custom_keywords():
    form = kb_form.SaveCustomKeywordsForm().validate_for_api()

    _words = list(set(form.words.data))
    _words.sort(key=form.words.data.index)

    words = []
    for word in _words:
        words.append(CustomKeywords(word=word, type=form.type.data, custom_id=str(form.uid.data)))

    CustomKeywords.objects.filter(custom_id=str(form.uid.data), type=form.type.data).delete()
    if words:
        CustomKeywords.objects.insert(words, load_bulk=False)
    return Success()


@api.route('/white_list_keywords', methods=['GET'])
def white_list_keywords():
    form = kb_form.QueryWhiteKeywordForm().validate_for_api()
    query = {'word__icontains': form.word.data} if form.word.data else {}
    query['custom_id'] = str(form.uid.data)
    page = WhiteListKeywords.objects.filter(**query).only('word').paginate(form.page.data, form.per_page.data)
    return jsonify({
        'items': list(page.items.as_pymongo()),
        'pageInfo': page.page_info()
    })


@api.route('/white_list_keywords', methods=['POST'])
def save_white_list_keywords():
    form = kb_form.SaveWhiteKeywordsForm().validate_for_api()

    _words = list(set(form.words.data))
    _words.sort(key=form.words.data.index)
    exist_words = WhiteListKeywords.objects.filter(word__in=_words, custom_id=str(form.uid.data)).scalar('word')

    words = []
    for word in _words:
        if word in exist_words:
            continue
        words.append(WhiteListKeywords(word=word, custom_id=str(form.uid.data)))

    # WhiteListKeywords.objects.filter(custom_id=str(form.uid.data)).delete()
    WhiteListKeywords.objects.insert(words, load_bulk=False)
    return Success()


@api.route('/delete_white_list_keywords', methods=['POST'])
def delete_white_list_keywords():
    form = kb_form.DeleteWhiteKeywordForm().validate_for_api()
    WhiteListKeywords.objects.filter(pk__in=form.ids.data, custom_id=str(form.uid.data)).delete()
    return Success()


@api.route('/get_profile', methods=['GET'])
def get_profile():
    form = kb_form.QueryProfileTemplateForm().validate_for_api()
    page = int(form.page.data) if form.page.data else 1
    limit = int(form.limit.data) if form.limit.data else 10
    uid = g.audit_uid
    if g.role.name == '管理员':
        uid = form.uid.data
    custom_query = {
        "product_name": "scanv",
        "is_disabled": False
    }
    system_query = {"is_disabled": {"$ne": True}}
    if name := form.name.data:
        custom_query['name__icontains'] = name
        system_query['display_name__zh_cn__icontains'] = name
    system_template = SystemProfileTemplate.objects.filter(**system_query)
    custom_template = CustomProfileTemplate.objects.filter(custom_id=str(uid), **custom_query)
    all_template = list(system_template)
    all_template.extend(custom_template)
    all_template.sort(key=lambda x: x.created_at, reverse=True)
    kb_vulns = get_kb_vuln()
    templates = []
    for template in all_template[max(page - 1, 0) * limit: page * limit]:
        templates.append(get_template_info(template, kb_vulns))

    data = {
        'data': {
            "max_count": limit,
            "page": page,
            "templates": templates,
            "total": len(all_template)
        },
        'message': 'ok',
        'code': 200,
    }

    return jsonify(data)


@api.route('/save_profile', methods=['POST'])
def save_profile():
    form = kb_form.SaveProfileTemplateForm().validate_for_api()
    name = form.name.data
    plugin_list = form.plugin_list.data
    uid = g.audit_uid
    if g.role.name == '管理员':
        uid = form.uid.data
    if check_custom_profile_name(str(uid), name):
        data = {
            "msg": f"名称 '{name}' 重复",
            "code": 403
        }
        return jsonify(data)
    if not plugin_list:
        data = {
            "msg": f"漏洞插件不能为空",
            "code": 403
        }
        return jsonify(data)
    description = form.description.data
    created_at = datetime.datetime.utcnow()
    updated_at = datetime.datetime.utcnow()
    custom_profile_template_id = uuid4().hex
    custom_profile_template = CustomProfileTemplate(
        name=name,
        plugin_list=plugin_list,
        description=description,
        created_at=created_at,
        updated_at=updated_at,
        custom_id=str(uid),
        product_name="scanv",
        custom_profile_template_id=custom_profile_template_id,
    )
    custom_profile_template.display_name = NameType(
        zh_cn=name
    )
    custom_profile_template.save()
    log_info = {
        'affects': [],
        'actionMessage': '新增用户自定义漏洞模板',
        'detail': f'新增用户自定义漏洞模板, 漏洞模板名称: {name}'
    }
    CProfileTemplateLog(**log_info).info(log_info['actionMessage'], log_info['detail'])
    return Success()


@api.route('/update_profile', methods=['POST'])
def update_profile():
    form = kb_form.UpdateProfileTemplateForm().validate_for_api()
    name = form.name.data
    uid = g.audit_uid
    if g.role.name == '管理员':
        uid = form.uid.data
    plugin_list = form.plugin_list.data
    description = form.description.data
    updated_at = datetime.datetime.utcnow()
    custom_profile_template_id = form.custom_profile_template_id.data
    query = {
        "custom_profile_template_id":custom_profile_template_id,
        "is_disabled": False,
        "product_name": "scanv"
    }
    custom_profile_template = CustomProfileTemplate.objects.filter(custom_id=str(uid), **query).first()
    if custom_profile_template.name != name and check_custom_profile_name(str(uid), name):
        data = {
            "msg": f"名称 '{name}' 重复",
            "code": 403
        }
        return jsonify(data)
    custom_profile_template.update(
        name=name,
        plugin_list=plugin_list,
        description=description,
        updated_at=updated_at,
        set__display_name__zh_cn=name,
    )

    log_info = {
        'affects': [],
        'actionMessage': '编辑用户自定义漏洞模板',
        'detail': f'编辑用户自定义漏洞模板, 漏洞模板名称: {name}'
    }
    UProfileTemplateLog(**log_info).info(log_info['actionMessage'], log_info['detail'])

    return Success()


@api.route('/deleted_profile', methods=['POST'])
def deleted_profile():
    form = kb_form.DeletedProfileTemplateForm().validate_for_api()
    uid = g.audit_uid
    if g.role.name == '管理员':
        uid = form.uid.data
    custom_profile_template_ids = form.custom_profile_template_ids.data
    query = {
        "custom_profile_template_id__in":custom_profile_template_ids,
        "is_disabled": False,
        "product_name": "scanv"
    }
    template_name_list = [item.name for item in CustomProfileTemplate.objects.filter(custom_id=str(uid), **query)]
    CustomProfileTemplate.objects.filter(custom_id=str(uid), **query).update(is_disabled=True)

    log_info = {
        'affects': [],
        'actionMessage': '删除用户自定义漏洞模板',
        'detail': f"删除漏洞模板 '{', '.join(template_name_list)}'"
    }
    CProfileTemplateLog(**log_info).info(log_info['actionMessage'], log_info['detail'])

    return Success()


@api.route('/custom_profile_count', methods=['GET'])
def get_custon_profile_count():
    form = kb_form.UidForm().validate_for_api()
    uid = g.audit_uid
    if g.role.name == '管理员':
        uid = form.uid.data
    query = {
        "is_disabled": False,
        "product_name": "scanv"
    }
    custom_profile_count = CustomProfileTemplate.objects.filter(custom_id=str(uid), **query).count()
    data = {
        "msg": "ok",
        "data": {
            "custom_profile_count": custom_profile_count,
        },
        "code": 200,
    }
    return jsonify(data)


@api.route('/get_custom_profile', methods=['POST'])
def get_custom_profile():
    form = kb_form.GetProfileTemplateForm().validate_for_api()
    custom_profile_template_id = form.custom_profile_template_id.data
    query = {
        "is_disabled": False,
        "product_name": "scanv"
    }
    custom_profile = CustomProfileTemplate.objects.filter(custom_profile_template_id=custom_profile_template_id, **query).first()
    if not custom_profile:
        data = {
            "mssage": "参数错误",
            "code": 403,
        }
        return jsonify(data)
    data = {
        "data": {
            "name": custom_profile.name,
            "description": custom_profile.description,
            "plugin_list": custom_profile.plugin_list,
            "custom_profile_template_id": custom_profile.custom_profile_template_id
        },
        "message": "ok",
        "code": 200
    }
    return jsonify(data)


@api.route('/get_all_profile_names', methods=['GET'])
def get_profile_names():
    form = kb_form.UidForm().validate_for_api()
    uid = g.audit_uid
    if g.role.name == '管理员':
        uid = form.uid.data
    system_query = {"is_disabled": {"$ne": True}}
    system_template = SystemProfileTemplate.objects.filter(**system_query)
    query = {
        "is_disabled": False,
        "product_name": "scanv"
    }
    custom_template = CustomProfileTemplate.objects.filter(custom_id=str(uid), **query)
    all_template = list(system_template)
    all_template.extend(custom_template)
    all_template.sort(key=lambda x: x.created_at, reverse=True)
    profile_names = []
    for template in all_template:
        profile_names.append(get_profile_name(template))
    data = {
        "data": profile_names,
        "massage": "ok",
        "code": 200
    }
    return jsonify(data)


def get_template_info(profile_template, kb_vulns):
    plugin_list = profile_template.plugin_list
    description = profile_template.description
    display_name_obj = profile_template.display_name
    plugins_count = len(plugin_list)
    display_name = {}
    plugin_severitys = {}
    plugin_types = {}
    if hasattr(profile_template, "custom_profile_template_id"):
        custom_profile_template_id = profile_template.custom_profile_template_id
    else:
        custom_profile_template_id = ""
    if hasattr(display_name_obj, "zh_cn"):
        display_name["zh_cn"] = display_name_obj.zh_cn
    if hasattr(display_name_obj, "en_us"):
        display_name["en_us"] = display_name_obj.en_us
    for plugin_key in plugin_list:
        if plugin_key not in kb_vulns:
            continue
        plugin_info = kb_vulns.get(plugin_key)
        plugin_severity = plugin_info.get("severity")
        plugin_name = plugin_info.get("name", {}).get("zh_cn")
        types = plugin_info.get("vuln_type")
        plugin_severitys.setdefault(plugin_severity, 0)
        plugin_severitys[plugin_severity] += 1
        for _type in types:
            type_name = _type.get("zh_cn")
            plugin_types.setdefault(type_name, [])
            plugin_types[type_name].append(plugin_name)

    template = {
        "custom_profile_template_id": custom_profile_template_id,
        "description": description,
        "display_name": display_name,
        "name": display_name.get("zh_cn"),
        "plugin_severity": plugin_severitys,
        "plugin_types": plugin_types,
        "plugins_count": plugins_count
    }
    return template


def get_kb_vuln():
    pipline = [
        {
            '$match': {
                'is_plugin': True,
                'is_deleted': False,
                'source': {'$in': ['kscan', 'nscan', 'cve', 'nuclei']}
            }
        }, {
            '$group': {
                '_id': '$source_id',
                'vuln_type': {
                    '$first': '$vuln_type'
                },
                'name': {
                    '$first': '$name'
                },
                'source_id': {
                    '$first': '$source_id'
                },
                'severity':{
                    '$first': '$severity'
                }
            }
        }
    ]
    vulns = {}
    for item in KbVuln.objects.aggregate(pipline):
        vulns[item['source_id']] = item
    return vulns


def check_custom_profile_name(uid, profile_name):
    query = {
        "is_disabled": False,
        "product_name": "scanv",
        "name": profile_name
    }
    system_query = {
        "is_disabled": {"$ne": True},
        "display_name__zh_cn": profile_name
    }
    is_custom_repeat_name = CustomProfileTemplate.objects.filter(custom_id=uid, **query).first()
    is_systom_repeat_name = SystemProfileTemplate.objects.filter(**system_query).first()
    if is_custom_repeat_name or is_systom_repeat_name:
        return True
    else:
        return False


def get_profile_name(profile_name):
    data = {
        "name": profile_name.name,
        "display_name": profile_name.display_name
    }
    if hasattr(profile_name, "value"):
        data["value"] = profile_name.value
        data["custom_profile_template_id"] = ""
    elif hasattr(profile_name, "custom_profile_template_id"):
        data["custom_profile_template_id"] = profile_name.custom_profile_template_id
        data["value"] = ""
    return data
